package main

import (
	"fmt"
)

func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
	totalLength := len(nums1) + len(nums2)
	// 中位数的奇数和偶数计算公式
	if totalLength%2 == 1 {
		midIndex := totalLength / 2
		return float64(getKthElement(nums1, nums2, midIndex+1))
	} else {
		midIndex1, midIndex2 := totalLength/2-1, totalLength/2
		return float64(getKthElement(nums1, nums2, midIndex1+1)+getKthElement(nums1, nums2, midIndex2+1)) / 2.0
	}
}

func getKthElement(nums1, nums2 []int, k int) int {
	index1, index2 := 0, 0
	for {
		// 数组1是空的
		if index1 == len(nums1) {
			// 直接取数组2的值
			return nums2[index2+k-1]
		}
		// 数组2是空的
		if index2 == len(nums2) {
			//直接取数组1的值
			return nums1[index1+k-1]
		}
		// k = 1的情况
		if k == 1 {
			return min(nums1[index1], nums2[index2])
		}
		half := k / 2
		// 取nums1[k/2 -1] nums2[k/2 -1]
		newIndex1 := min(index1+half, len(nums1)) - 1
		newIndex2 := min(index2+half, len(nums2)) - 1
		pivot1, pivot2 := nums1[newIndex1], nums2[newIndex2]
		// 根据结果更新k的值和下次循环的指针下标
		if pivot1 <= pivot2 {
			k -= (newIndex1 - index1 + 1)
			index1 = newIndex1 + 1
		} else {
			k -= (newIndex2 - index2 + 1)
			index2 = newIndex2 + 1
		}
	}
}

func min(x, y int) int {
	if x < y {
		return x
	}
	return y
}

func main() {
	nums1 := []int{1, 2, 3}
	nums2 := []int{3, 3, 3}

	res1 := findMedianSortedArrays(nums1, nums2)
	fmt.Println(res1)

	nums3 := []int{1, 2}
	nums4 := []int{3, 4}
	res2 := findMedianSortedArrays(nums3, nums4)
	fmt.Println(res2)
}
